package com.yupi.yuaiagent.agent;


import com.yupi.yuaiagent.agent.model.AgentState;
import io.reactivex.Emitter;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

@Slf4j
@Data
public abstract class BaseAgent {

    // 助手名称
    private String name;

    // 提示词
    private String systemPrompt;
    private String nextStepPrompt;

    // 运行状态 （空闲中）
    private AgentState state = AgentState.IDLE;

    // 执行步骤
    private int currentStep;  // 当前执行步骤
    private int maxSteps;   // 最大执行步骤10步

    // 大模型
    private ChatClient  chatClient;

    // Memory 自己维护上下文
    List<Message> messages = new ArrayList<>();

    /**
     * 运行
     * @return
     */
    public String run(String userPrompt){
        // 检查状态
        if(this.state != AgentState.IDLE){
            throw new RuntimeException("助手正在运行中");
        }
        // 检测用户提示词是否为空
        if(userPrompt.isBlank()){
            throw new RuntimeException("用户输入为空");
        }
        // 更改状态 修改状态为运行中
        this.state = AgentState.RUNNING;
        // 记录消息上下文
        messages.add(new UserMessage(userPrompt));
        // 结果保存列表
        ArrayList<String> results = new ArrayList<>();
        try {
            for(int i = 0; i < maxSteps && state != AgentState.FINISHED; i++){
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("当前执行步骤：" + currentStep + "/" + maxSteps);
                // 单步执行
                String stepResult = step();
                String result = "Step：" + currentStep + ":" + stepResult;
                results.add(result);
            }
            if(currentStep >= maxSteps){
                this.state  = AgentState.FINISHED;
                results.add("Terminated：Reach max steps (" + maxSteps + ")");
            }
            // 返回结果
            return String.join("\n", results);
        } catch (Exception e) {
            this.state = AgentState.ERROR;
            log.error("Error occurred during agent running: " + e.getMessage());
            return "执行错误: " + e.getMessage();
        } finally {
            // 清理资源
            this.clean();
        }
    }


    /**
     * 运行代理 (流式输出)
     * @param userPrompt 用户提示词
     * @return SseEmitter实例
     */
    public SseEmitter runStream(String userPrompt){
        SseEmitter emitter = new SseEmitter(10000L);

        // 异步执行。
        CompletableFuture.runAsync(() -> {
            try {
                // 检查状态
                if(this.state != AgentState.IDLE){
                    emitter.send("智能体状态错误,此时不处于空闲状态");
                    emitter.complete();
                    return;
                }
                // 检测用户提示词是否为空
                if(userPrompt.isBlank()){
                    emitter.send("智能体不支持用户的提示词为空");
                    emitter.complete();
                    return;
                }

                // 更改状态 修改状态为运行中
                this.state = AgentState.RUNNING;
                // 记录消息上下文
                messages.add(new UserMessage(userPrompt));

                try {

                    for(int i = 0; i < maxSteps && state != AgentState.FINISHED; i++){
                        int stepNumber = i + 1;
                        currentStep = stepNumber;
                        log.info("当前执行步骤：" + currentStep + "/" + maxSteps);
                        // 单步执行
                        String stepResult = step();
                        String result = "Step：" + currentStep + ":" + stepResult;
                        // 发送每一步结果
                        emitter.send(result);
                    }

                    // 检查是否超出最大执行步骤
                    if(currentStep >= maxSteps){
                        this.state  = AgentState.FINISHED;
                        emitter.send("执行结束：达到最大步骤 (" + maxSteps + ")");
                    }

                    // 正常完成
                    emitter.complete();
                } catch (Exception e) {
                    this.state = AgentState.ERROR;
                    log.error("Error occurred during agent running: " + e.getMessage());
                    try {
                        emitter.send("智能体执行错误: " + e.getMessage());
                        emitter.complete();
                    } catch (IOException ex) {
                        emitter.completeWithError(ex);
                    }
                } finally {
                    // 清理资源
                    this.clean();
                }
            } catch (IOException e) {
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }

    /**
     * 单步执行步骤
     * @return
     */
    public abstract String step();


    /**
     * 清理资源 子类可以重写这个类清理上下文
     */
    public abstract void clean();
}
